import Foundation
import imp
import objc
import os
import shutil
import sys
import StringIO

def sampleRateForRecords( records ):
    sr = -1
    for r in records:
        if sr == -1:
            sr = r.sampleRate()
        if sr != r.sampleRate():
            return -1
    return sr
    
def getFFTSize( records ):
    fftSize = 0
    for r in records:
        fftSize = max( fftSize, r.fftLength() )
    return fftSize

class RMSAverage (objc.lookUpClass('SMUGFMPlugIn')):
    _windowController = None
    
    def plugInDidLoad(self):
        self.plugInManager().addPlugInMenuItemNamed_withParentMenuTitle_target_action_keyEquivalent_keyEquivalentModifierMask_(
            "RMS Average",
            None,
            self,
            'handleRunScript:userObject:',
            '',
            0)

    def handleRunScript_userObject_(self, windowController, userObject):
        document = windowController.document()
        records = windowController.selectedMeasurementRecords();
        recordCount = records.count()

        if recordCount == 0:
            windowController.scriptError_( "Cannot perform RMS Average - no records selected" )
        else:
            sampleRate = sampleRateForRecords( records )

            if sampleRate == -1:
                windowController.scriptError_( "Cannot perform RMS Average - mismatched sample rates" )
            else:
                N = int(getFFTSize( records ))

                rmsArray = []
                for r in records:
                    wir = r.windowedImpulseResponse()
                    wir.setLength_(N)
                    magResponse = wir.fftWithLength_(N)
                    rmsArray.append(magResponse)

                complexVec = objc.lookUpClass("SMUGComplexVector").RMSAverageOfVectors_( rmsArray )
                ir = complexVec.ifft()

                newRecord = document.addNewRecordWithImpulse_sampleRate_(ir, sampleRate)
                newRecord.setComment_( "RMS Average" )
                newRecord.expandWindow()